import torch
import os
import numpy as np
from model.mesh_dataset import MeshDataset
from model.model_triangle_generator import ModelTriangleGenerator
import re  # 引入正则表达式库

# 定义项目根路径
PROJECT_ROOT = r"C:\Users\47306\PycharmProjects\Mesh-simplification-GNN"

# 配置环境变量
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

def save_obj(vertices, faces, save_path):
    """
    将生成的网格保存为 .obj 文件
    """
    with open(save_path, 'w') as obj_file:
        for v in vertices:
            obj_file.write(f"v {v[0]} {v[1]} {v[2]}\n")  # 保存顶点信息
        for f in faces:
            obj_file.write(f"f {f[0]+1} {f[1]+1} {f[2]+1}\n")  # 保存面信息 (OBJ从1开始计数)
    print(f"Saved OBJ file to {save_path}")

def viz_model_train_and_save_obj():
    # 设置训练参数
    number_neigh_tri = 20
    torch_dataset = MeshDataset(os.path.join(PROJECT_ROOT, "3d_models/stl/"))  # 使用全局路径拼接
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 获取保存的模型文件名
    saved_model_filenames = [f for f in os.listdir(os.path.join(PROJECT_ROOT, 'save_models/triangle_generator/1')) if
                             f.endswith('.pt')]  # 只选择pt文件
    print(f"Found {len(saved_model_filenames)} saved models")

    for save_model in saved_model_filenames:
        model_path = os.path.join(PROJECT_ROOT, 'save_models/triangle_generator/1', save_model)  # 拼接模型文件路径
        gnn_model = ModelTriangleGenerator(number_neigh_tri).to(device)
        gnn_model.load_state_dict(torch.load(model_path))  # 加载模型
        gnn_model.eval()

        with torch.no_grad():
            # 使用正则表达式提取文件名中的数字部分
            match = re.search(r'(\d+)', save_model)  # 提取连续的数字部分
            if match:
                index_img = int(match.group(0))  # 获取匹配到的数字并转换为整数
            else:
                print(f"Warning: Model filename {save_model} does not contain a valid number.")
                continue  # 如果文件名不符合预期，跳过该文件

            # 生成索引，用于设置不同视角的角度
            angle = index_img % 360
            torch_graph, _ = torch_dataset[0]  # 读取数据集中的第一个图形数据

            # 简化率 默认为0.8
            target_number_point = int(
                max(200, min(len(torch_graph.x) * 0.8, 25_000)))
            # 通过模型生成三角形网格
            triangles = gnn_model(target_number_point, torch_graph.to(device))  # 假设每次生成200个三角形
            vertices = triangles.reshape(-1, 3).cpu().numpy()  # 将输出重塑为顶点坐标

            # 去重：合并重复的顶点，并重新索引
            unique_vertices, unique_indices = np.unique(vertices, axis=0, return_inverse=True)
            num_vertices = unique_vertices.shape[0]

            # 创建面（face）索引，并使用去重后的顶点索引
            num_triangles = triangles.shape[0]
            faces = np.arange(num_triangles * 3).reshape(-1, 3)

            # 确保面是由去重后的顶点索引组成
            faces = unique_indices[faces]

            # 中心化并缩放网格
            centroid = np.mean(unique_vertices, axis=0)
            unique_vertices -= centroid  # 中心化网格
            max_distance = np.max(np.linalg.norm(unique_vertices, axis=1))
            unique_vertices /= max_distance  # 缩放网格，使得最大距离为1

            # 生成.obj文件保存路径
            save_dir = os.path.join(PROJECT_ROOT, "saved_obj", "model_viz", "1")
            os.makedirs(save_dir, exist_ok=True)  # 如果文件夹不存在，则创建
            save_path = os.path.join(save_dir, f"{index_img}.obj")  # OBJ文件保存路径

            # 保存为.obj文件
            save_obj(unique_vertices, faces, save_path)

# 调用函数生成并保存 .obj 文件
viz_model_train_and_save_obj()
